package com.xfsg.core.rocketmq.hook;

import com.yomahub.tlog.constant.TLogConstants;
import com.yomahub.tlog.core.rpc.TLogLabelBean;
import com.yomahub.tlog.core.rpc.TLogRPCHandler;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.client.hook.ConsumeMessageContext;
import org.apache.rocketmq.client.hook.ConsumeMessageHook;
import org.apache.rocketmq.common.message.MessageExt;

/**
 * @author ZhangYi
 * @classname ConsumerRequestIdtHook
 * @description TODO
 * @date 2022/2/18 22:58
 */
@Slf4j
public class ConsumerLogTraceHook implements ConsumeMessageHook {

    private static TLogRPCHandler tLogRPCHandler = new TLogRPCHandler();

    @Override
    public String hookName() {
        return "ConsumerLogTraceHook";
    }

    @Override
    public void consumeMessageBefore(ConsumeMessageContext context) {
        if (context == null || context.getMsgList() == null || context.getMsgList().isEmpty()) {
            return;
        }
        if (context.getMsgList().size() > 1) {
            log.warn("greater 1 message, consumerGroup={}", context.getConsumerGroup());
        }
        for (MessageExt messageExt : context.getMsgList()) {
            if (messageExt == null) {
                continue;
            }
            String traceId = messageExt.getUserProperty(TLogConstants.TLOG_TRACE_KEY);
            String preIvkApp = messageExt.getUserProperty(TLogConstants.PRE_IVK_APP_KEY);
            String preIvkHost = messageExt.getUserProperty(TLogConstants.PRE_IVK_APP_HOST);
            String preIp = messageExt.getUserProperty(TLogConstants.PRE_IP_KEY);
            String spanId = messageExt.getUserProperty(TLogConstants.TLOG_SPANID_KEY);
            TLogLabelBean labelBean = new TLogLabelBean(preIvkApp, preIvkHost, preIp, traceId, spanId);
            tLogRPCHandler.processProviderSide(labelBean);
        }
    }

    @Override
    public void consumeMessageAfter(ConsumeMessageContext context) {
        if (context == null || context.getMsgList() == null || context.getMsgList().isEmpty()) {
            return;
        }
        tLogRPCHandler.cleanThreadLocal();
    }
}
